import torch
from torch import nn
from torchvision import transforms as T
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader, random_split
import pytorch_lightning as pl
from torchmetrics import Accuracy


class Backbone(torch.nn.Module):
    def __init__(self, hidden_dim=128):
        super().__init__()
        self.l1 = torch.nn.Linear(28 * 28, hidden_dim)
        self.l2 = torch.nn.Linear(hidden_dim, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = torch.relu(self.l1(x))
        x = torch.relu(self.l2(x))
        return x

net = nn.Sequential(
    nn.Conv2d(in_channels=1 ,out_channels=32 ,kernel_size = 3),
    nn.MaxPool2d(kernel_size = 2 ,stride = 2),
    nn.Conv2d(in_channels=32 ,out_channels=64 ,kernel_size = 5),
    nn.MaxPool2d(kernel_size = 2 ,stride = 2),
    nn.Dropout2d(p = 0.1),
    nn.AdaptiveMaxPool2d((1 ,1)),
    nn.Flatten(),
    nn.Linear(64 ,32),
    nn.ReLU(),
    nn.Linear(32 ,10)
)

class Model(pl.LightningModule):

    def __init__(self ,net):
        super().__init__()
        self.save_hyperparameters(ignore=["backbone"])
        self.net = net
        self.train_acc = Accuracy(task="multiclass", num_classes = 10, top_k=2)
        self.val_acc = Accuracy(task="multiclass", num_classes = 10, top_k=2)
        self.test_acc = Accuracy(task="multiclass", num_classes = 10, top_k=2)

    def forward(self ,x):
        x = self.net(x)
        return x

    # 定义loss
    def training_step(self, batch, batch_idx):
        x, y = batch
        preds = self(x)
        loss = nn.CrossEntropyLoss()(preds ,y)
        self.logger.log_metrics({"training_step_loss": loss}, step=batch_idx)
        # print("training_step ---------")
        return {"loss" :loss ,"preds" :preds.detach() ,"y" :y.detach()}

    def training_step_end(self ,outputs):
        train_acc = self.train_acc(outputs['preds'], outputs['y']).item()
        self.logger.log_metrics({"training_step_end_acc": train_acc})
        # print("training_step_end +++++++++++++")
        return {"loss" :outputs["loss"].mean()}

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.forward(x)
        loss = nn.CrossEntropyLoss()(y_hat ,y)
        # print("validation_step ---------")
        self.logger.log_metrics({"validation_step_loss": loss}, step=batch_idx)
        return {"loss" :loss ,"y_hat" :y_hat.detach() ,"y" :y.detach()}

    def validation_step_end(self ,outputs):
        val_acc = self.val_acc(outputs['y_hat'], outputs['y']).item()
        # print("validation_step_end +++++++++++")
        self.logger.log_metrics({"validation_step_end_acc": val_acc})
        self.logger.log_metrics({"validation_step_end_loss": outputs["loss"].mean() })

    def test_step(self, batch, batch_idx):
        x, y = batch
        preds = self(x)
        loss = nn.CrossEntropyLoss()(preds ,y)
        self.logger.log_metrics({"test_step_loss": loss}, step=batch_idx)
        return {"loss" :loss ,"preds" :preds.detach() ,"y" :y.detach()}

    def test_step_end(self ,outputs):
        test_acc = self.test_acc(outputs['preds'], outputs['y']).item()
        self.logger.log_metrics({"test_step_end_acc": test_acc})
        self.logger.log_metrics({"test_step_end_loss": outputs["loss"].mean()})

    # 定义optimizer,以及可选的lr_scheduler
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)

if __name__ == '__main__':

    model = Model(net)

    # 查看模型大小
    model_size = pl.utilities.memory.get_model_size_mb(model)
    print("model_size = {} M \n".format(model_size))
    # model.example_input_array = [features]
    summary = pl.utilities.model_summary.ModelSummary(model ,max_depth=-1)
    print(summary)